"""
DSPy豆包模型集成解决方案
日期：2025年6月21日
"""
import time
import dspy
from typing import List, Optional, Any
from dspy.teleprompt import BootstrapFewShot

from doubaoV2 import DoubaoLM


class BasicQA(dspy.Signature):
    """回答用户的问题"""
    question = dspy.InputField(desc="用户提出的问题")
    answer = dspy.OutputField(desc="对问题的清晰、准确回答")


class QA(dspy.Module):
    """问答模块 - 完全自定义实现"""

    def __init__(self):
        super().__init__()
        self.predictor = dspy.Predict(BasicQA)

    def forward(self, question: str) -> dspy.Prediction:
        """处理用户问题并返回答案"""
        return self.predictor(question=question)


def create_dataset() -> List[dspy.Example]:
    """创建问答数据集
    返回:
        List[dspy.Example]: 包含问答示例的列表
    """
    return [
        dspy.Example(question="中国的首都在哪里？", answer="北京").with_inputs("question"),
        dspy.Example(question="Python是什么？", answer="Python是一种高级编程语言").with_inputs("question"),
        dspy.Example(question="太阳系有多少颗行星？", answer="太阳系有八颗行星").with_inputs("question"),
        dspy.Example(question="水的化学式是什么？", answer="H₂O").with_inputs("question"),
        dspy.Example(question="法国的首都是哪里？", answer="巴黎").with_inputs("question"),
    ]


def evaluate_qa(example: dspy.Example, prediction: dspy.Prediction, trace: Optional[Any] = None) -> bool:
    """评估函数 - 检查预测答案是否匹配真实答案
    参数:
        example (dspy.Example): 包含真实答案的示例对象
        prediction (dspy.Prediction): 模型预测结果对象
        trace (Optional[Any]): 可选的调试跟踪信息
    返回:
        bool: 预测是否正确
    """
    # 确保预测对象包含answer属性
    if not hasattr(prediction, 'answer'):
        print(f"预测对象缺少answer属性: {type(prediction)}")
        return False

    # 获取真实答案和预测答案
    ground_truth = example.answer.lower().strip()
    predicted = prediction.answer.lower().strip()

    print(f"评估: 问题='{example.question}'")
    print(f"真实答案: '{ground_truth}'")
    print(f"预测答案: '{predicted}'")

    # 简化的评估逻辑 - 真实答案是否包含在预测答案中
    is_correct = ground_truth in predicted
    print(f"评估结果: {'正确' if is_correct else '错误'}")
    return is_correct


def compile_model(train_set, lm: DoubaoLM) -> QA:
    """使用引导式优化提升模型性能
    参数:
        train_set (List[dspy.Example]): 训练数据集
        lm (DoubaoLM): 豆包语言模型实例
    返回:
        QA: 优化后的问答模型
    """
    # 创建基础QA模块
    qa_model = QA()

    # 配置DSPy使用豆包模型
    dspy.configure(lm=lm)

    # BootstrapFewShot 的核心目标是优化提示词 (prompt)，而不是直接优化模型参数或权重
    # 主要作用是通过自动化提示工程（Automated Prompt Engineering）来提高语言模型的性能
    # 核心创新在于将提示工程从一门艺术转变为可编程、自动化的过程。
    # 通过系统地探索提示空间，自动发现最优提示配置，显著提升了语言模型在特定任务上的表现。
    teleprompter = BootstrapFewShot(metric=evaluate_qa)

    print("开始优化模型...")
    optimized_model = teleprompter.compile(qa_model, trainset=train_set)
    print("模型优化完成！")

    # 打印最终优化后的提示
    print("\n" + "="*50)
    print("优化后的提示模板:")
    print("="*50)
    print(optimized_model.predictor.signature.instructions)
    print("="*50)

    return optimized_model

def test_model(model: QA, questions: List[str]) -> None:
    """测试模型性能
    参数:
        model (QA): 问答模型实例
        questions (List[str]): 测试问题列表
    """
    print("\n测试问答模型：")
    for question in questions:
        # 调用模型获取响应
        prediction = model(question=question)

        # 确保我们有answer属性
        if hasattr(prediction, 'answer'):
            response = prediction.answer
        else:
            response = str(prediction)
            print(f"预测对象缺少answer属性: {type(prediction)}")

        print(f"问题: {question}")
        print(f"回答: {response}")
        print("-" * 50)

def interactive_mode(model: QA, lm: DoubaoLM) -> None:
    """用户交互模式
    参数:
        model (QA): 问答模型实例
        lm (DoubaoLM): 豆包语言模型实例
    """
    print("\n交互模式：输入问题或输入'退出'结束")
    while True:
        try:
            user_input = input("\n你的问题: ")
            if user_input.lower() in ["退出", "exit", "quit"]:
                break

            # 调用模型获取回答
            start_time = time.time()
            prediction = model(question=user_input)
            response_time = time.time() - start_time

            # 确保我们有answer属性
            if hasattr(prediction, 'answer'):
                response = prediction.answer
            else:
                response = str(prediction)
                print(f"预测对象缺少answer属性: {type(prediction)}")

            print(f"\n智能体回答: {response}")
            print(f"响应时间: {response_time:.2f}秒")
        except KeyboardInterrupt:
            print("\n程序已终止")
            break
        except Exception as e:
            print(f"处理问题时出错: {str(e)}")


def main(api_key: str):
    # 初始化豆包模型
    lm = DoubaoLM(
        api_key=api_key,
    )

    # 创建数据集
    train_set = create_dataset()

    # 编译优化模型
    optimized_qa = compile_model(train_set, lm)

    # 测试问题
    test_questions = [
        "中国的首都在哪里？",
        "Python的主要特点是什么？",
        "太阳系最大的行星是哪颗？",
        "水的沸点是多少摄氏度？",
    ]

    # 测试模型
    test_model(optimized_qa, test_questions)

    # 进入用户交互模式
    interactive_mode(optimized_qa, lm)


if __name__ == "__main__":
    import fire
    fire.Fire(main)
